#include "buddy2.h"
#include "util.h"

#define LEFT_LEAF(index)        ((index)*2+1)
#define RIGHT_LEAF(index)       ((index)*2+2)
#define PARENT(index)           (((index)+1)/2-1)

#define IS_POWER_OF_2(x)        (!((x) & ((x) -1)))
#define MAX(a,b)                ((a) > (b) ? (a):(b))

#define ALLOC                   util_malloc
#define FREE                    util_free

struct buddy2
{
    unsigned size;
    unsigned longest[1];
};

static unsigned fix_size(unsigned size)
{
    size |=  size >> 1;
    size |=  size >> 2;
    size |=  size >> 4;
    size |=  size >> 8;
    size |=  size >> 16;

    return size + 1;
}

struct buddy2 *buddy2_new(int size)
{
    struct buddy2 *self = NULL;
    unsigned node_size;
    int i;

    if(size < 1 || !IS_POWER_OF_2(size))
        return NULL;

    self = (struct buddy2 *)ALLOC(2 * size * sizeof(unsigned));
    self->size = size;
    node_size = size * 2;

    for(i = 0; i < 2 * size - 1; ++i)
    {
        if(IS_POWER_OF_2(i+1))
            node_size /= 2;

        self->longest[i] = node_size;
    }

    return self;
}

void buddy2_destroy(struct buddy2 *self)
{
    FREE(self);
}

int buddy2_alloc(struct buddy2 *self, int size)
{
    unsigned node_size;
    unsigned index = 0, offset = 0;

    if(self == NULL)
        return -1;

    if(size <= 0)
        size = 1;
    else if(!IS_POWER_OF_2(size))
        size = fix_size(size);

    if(self->longest[index] < size)
        return -1;

    for(node_size = self->size; node_size != size; node_size /= 2)
    {
        if(self->longest[LEFT_LEAF(index)] >= size)
            index = LEFT_LEAF(index);
        else
            index = RIGHT_LEAF(index);
    }

    self->longest[index] = 0;

    offset = (index + 1) * node_size - self->size;

    while(index)
    {
        index = PARENT(index);
        self->longest[index] = 
            MAX(self->longest[LEFT_LEAF(index)], self->longest[RIGHT_LEAF(index)]);
    }
    return offset;
}

void buddy2_free(struct buddy2 *self, int offset)
{
    unsigned node_size, index = 0;
    unsigned left_longest, right_longest;


    UTIL_ASSERT(self && offset >= 0 && offset < self->size);

    node_size = 1;
    index = offset + self->size - 1;

    for(; self->longest[index]; index = PARENT(index))
    {
        node_size *= 2;
        if(index == 0)
            return;
    }

    self->longest[index] = node_size;

    while(index)
    {
        index = PARENT(index);
        node_size *= 2;

        left_longest = self->longest[LEFT_LEAF(index)];
        right_longest = self->longest[RIGHT_LEAF(index)];

        if(left_longest + right_longest == node_size)
            self->longest[index] = node_size;
        else
            self->longest[index] = MAX(left_longest, right_longest);
    }
}

int buddy2_size(struct buddy2 *self, int offset)
{
    unsigned node_size, index = 0;

    UTIL_ASSERT(self && offset >= 0 && offset < self->size);

    node_size = 1;

    for(index = offset + self->size - 1; self->longest[index]; index = PARENT(index))
    {
        node_size *= 2;
    }

    return node_size;
}

unsigned buddy2_remaining_size(struct buddy2 *self)
{
    unsigned node_size;
    unsigned total_size = 0;
    unsigned index, j;
#if 0
    for(index = self->size * 2 - 1 - self->size;
            index < self->size * 2 - 1; index++)
#else
    for(index = self->size - 1; index < self->size * 2 - 1; index++)
#endif
    {
        j = index;

        while(j != 0)
        {
            node_size = 1;
            if(self->longest[j] == 0)
            {
                node_size = 0;
                break;
            }
            j = PARENT(j);
        }

        total_size += node_size;
    }

    
    return total_size;
}

void buddy2_dump_tree(struct buddy2 *self)
{
    unsigned node_num, ele_n_line;
    unsigned i, prev = 1, space_unit = 2;

    node_num = self->size * 2 - 1;

    ele_n_line = self->size;

    for(;;)
    {
        if(ele_n_line == 0 || node_num == 0)
            break;

        for(i = node_num - ele_n_line; i < node_num; i++)
        {
            if(i == node_num - ele_n_line)
                UTIL_PRINTF("%*d", prev, self->longest[i]);
            else
                UTIL_PRINTF("%*d", space_unit, self->longest[i]);
        }

        space_unit += prev * 2;
        prev = space_unit / 2;
        //UTIL_PRINTF(" -----------%d--%d\n", ele_n_line, node_num);
        UTIL_PRINTF("\r\n\r\n");

        node_num -= ele_n_line;
        ele_n_line /= 2;
    }

}

int buddy2_test1(void)
{
    struct buddy2 *self = NULL;
    int mem[] = {3, 8, 1, 32, 1, 2, 3, 2};
    int offset[] = {-1,-1,-1,-1,-1,-1,-1,-1};
    int n = sizeof(offset)/sizeof(offset[0]);
    int i;

    UTIL_ASSERT(n == sizeof(mem)/sizeof(mem[0]));

    self = buddy2_new(64);
    
    buddy2_dump_tree(self);

    UTIL_PRINTF("*****************************\r\n");
    for(i = 0; i < n; ++i)
    {
        offset[i] = buddy2_alloc(self, mem[i]);
        UTIL_PRINTF("=====alloc sz:%d =========\r\n", mem[i]);
        buddy2_dump_tree(self);

    }
    for(i = 0; i < n; ++i)
    {
        buddy2_free(self, offset[i]);
        UTIL_PRINTF("=====free of:%d ==========\r\n", offset[i]);
        buddy2_dump_tree(self);
    }

    buddy2_destroy(self);

    return 0;

}

int buddy2_test_main(void)
{
    return buddy2_test1();
}

